import os
import torch
from torch.utils.data import DataLoader
import numpy as np
from utils.param import Param

def get_loader(param: Param):
    datas = np.load(os.path.join(param.data_path, param.file_name), allow_pickle=True)
    data, word2ix, ix2word = datas['data'], datas['word2ix'], datas['ix2word']
    data = torch.from_numpy(data)
    train_loader = DataLoader(data, batch_size=param.batch_size, shuffle=True)
    # print(len(train_loader))
    return train_loader


def get_word2ix(param: Param):
    datas = np.load(os.path.join(param.data_path, param.file_name), allow_pickle=True)
    word2ix = datas['word2ix'].item()
    return word2ix

def get_ix2word(param:Param):
    datas = np.load(os.path.join(param.data_path, param.file_name), allow_pickle=True)
    ix2word = datas['ix2word'].item()
    return ix2word

if __name__ == '__main__':
    param = Param()
    trainloader = get_loader(param)
    for data in trainloader:
        print(data)
        break
    print(trainloader)

